import numpy as np
import cv2
import os
import json

class Calibration(object):
    def __init__(self, filepath):

        self.base_dir = os.path.dirname(__file__)

        self.visualization = False
        if self.visualization:
            self.v_dir = os.path.join(self.base_dir, "visualization")
            if not os.path.exists(self.v_dir):
                os.mkdir(self.v_dir)

        # 标定公共参数配置
        self.imageFormat = ".bmp"
        self.boardSize = (11, 8)
        self.squareSize = 15  # mm
        self.img_num = 22  # 根据图像名称命名规则为L_1.bmp, L_2.bmp ...; R_1.bmp, R_2.bmp, ...
        self.imageSize = (2448,2048)

        # prepare object points, like (0,0,0), (10,0,0), (20,0,0) ....,(100,70,0)
        self.objp = np.zeros((self.boardSize[0]*self.boardSize[1], 3), np.float32)
        self.objp[:, :2] = np.mgrid[0:self.boardSize[0], 0:self.boardSize[1]].T.reshape(-1, 2)
        self.world_points = self.objp * self.squareSize
        
        # self.monocular_calibrate(filepath, camera="R")

        # 双目标定参数配置
        # self.objectPoints = [self.world_points for _ in range(self.img_num)]
        self.stereo_calibrate(filepath)

    def monocular_calibrate(self, cal_path, camera_view):

        worldPoints = []
        imagePoints = []

        imagePathList = [os.path.join(cal_path, camera_view + "_" + str(i) + self.imageFormat) for i in range(1, self.img_num+1)]
        # print(imagePathList)
        for i in range(len(imagePathList)):

            image_bgr = cv2.imread(imagePathList[i])
            
            image_gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)

            ret, corners = cv2.findChessboardCorners(image_gray, self.boardSize, None, cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE)
            
            if ret:
                new_corners = cv2.cornerSubPix(image_gray, corners, (11, 11), (-1, -1), (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 300, 0.001))

                if self.visualization:
                    image_with_corners = cv2.drawChessboardCorners(image_bgr, self.boardSize, new_corners, ret)
                    # 显示带有角点的图像
                    cv2.namedWindow("corners",0);
                    cv2.resizeWindow("corners", 640, 480);
                    cv2.imshow('corners', image_with_corners)
                    cv2.waitKey(1000)  # 等待按键按下
                    cv2.destroyAllWindows()  # 关闭所有窗口
   
                worldPoints.append(self.world_points)
                imagePoints.append(new_corners)

        rms, CameraMat, DistCoeff, rvecs, tvecs = cv2.calibrateCamera(worldPoints, imagePoints, self.imageSize, None, None)

        # 将标定参数写入json文件
        rvecs_list = [rvec.tolist() for rvec in rvecs]
        tvecs_list = [tvec.tolist() for tvec in tvecs]
        worldPoints_list = [p.tolist() for p in worldPoints]
        imagePoints_list = [p.tolist() for p in imagePoints]
        calibration_params = {  
            'rms': rms,
            'CameraMat': CameraMat.tolist(),
            'DistCoeff': DistCoeff.tolist(),
            'rvecs': rvecs_list,
            'tvecs': tvecs_list,
            'worldPoints': worldPoints_list,
            'imagePoints': imagePoints_list,
            'imagePaths': imagePathList,
        }
        # 写入到文件  
        with open('monocular_calibiation_' + camera_view + '12101.json', 'w') as f:  
            json.dump(calibration_params, f, indent=4)  # indent参数用于美化输出，增加可读性

        # 反投影误差
        total_error = 0
        for i in range(len(worldPoints)):
            imagePoints_project, _ = cv2.projectPoints(worldPoints[i], rvecs[i], tvecs[i], CameraMat, DistCoeff)
            error = cv2.norm(imagePoints[i], imagePoints_project, cv2.NORM_L2) / len(imagePoints_project)
            total_error += error
        print("total error: ", total_error / len(worldPoints))

        return CameraMat, DistCoeff, rvecs, tvecs, imagePoints

    def stereo_calibrate(self, cal_path):

        CameraMatL, DistCoeffL, rvecsL, tvecsL, imagePointsL = self.monocular_calibrate(cal_path, camera_view="L")
        CameraMatR, DistCoeffR, rvecsR, tvecsR, imagePointsR = self.monocular_calibrate(cal_path, camera_view="R")

        assert len(imagePointsL)==len(imagePointsR), "The number of images in the left and right views is not equal!"

        objectPoints = [self.world_points for _ in range(len(imagePointsL))]

        # flags = 0
        # flags |= cv2.CALIB_FIX_INTRINSIC
        # flags |= cv2.CALIB_FIX_PRINCIPAL_POINT
        # flags |= cv2.CALIB_USE_INTRINSIC_GUESS
        # flags |= cv2.CALIB_FIX_FOCAL_LENGTH
        # flags |= cv2.CALIB_FIX_ASPECT_RATIO
        # flags |= cv2.CALIB_ZERO_TANGENT_DIST
        # flags |= cv2.CALIB_RATIONAL_MODEL
        # flags |= cv2.CALIB_SAME_FOCAL_LENGTH
        # flags |= cv2.CALIB_FIX_K3
        # flags |= cv2.CALIB_FIX_K4
        # flags |= cv2.CALIB_FIX_K5

        stereocalib_criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 100, 1e-5)
        
        rms, newCameraMatL, newDistCoeffL, newCameraMatR, newDistCoeffR, R, T, E, F = cv2.stereoCalibrate(
            objectPoints, imagePointsL, imagePointsR,
            CameraMatL, DistCoeffL,
            CameraMatR, DistCoeffR,
            self.imageSize,
            criteria=stereocalib_criteria,
            # flags=cv2.CALIB_FIX_INTRINSIC
            )


        # 将双目标定参数写入json文件
        # R_list = [v.tolist() for v in R]
        # T_list = [v.tolist() for v in T]
        # print(np.array(v_rot_list).shape, np.array(v_trans_list).shape)
        calibration_params = {
            'rms': rms,
            'CameraMatL': newCameraMatL.tolist(),
            'DistCoeffL': newDistCoeffL.tolist(),
            'CameraMatR': newCameraMatR.tolist(),
            'DistCoeffR': newDistCoeffR.tolist(),
            'R': R.tolist(),
            'T': T.tolist(),
            'E': E.tolist(),
            'F': F.tolist()
        }
  
        # 写入到文件  
        with open('stereo_calibiate12101.json', 'w') as f:  
            json.dump(calibration_params, f, indent=4)  # indent参数用于美化输出，增加可读性


class CalibrationParams(object):
    def __init__(self):
        pass

    
    def load_monocular_param(self, param_path):
        with open(param_path, 'r') as f:
            json_dict = json.load(f)

        self.CameraMat = np.array(json_dict['CameraMat']) # 3*3
        self.DistCoeff = np.array(json_dict['DistCoeff']) # 1*5
        self.rvecs = np.array(json_dict['rvecs']) # 38*3*1
        self.tvecs = np.array(json_dict['tvecs']) # 38*3*1
        self.worldPoints = np.array(json_dict['worldPoints']) # 38*88*3
        self.imagePoints = np.array(json_dict['imagePoints']) # 38*88*1*2
        self.imagePaths = json_dict['imagePaths']

    def load_stereo_param(self, param_path):
        with open(param_path, 'r') as f:
            json_dict = json.load(f)

        self.CameraMatL = np.array(json_dict['CameraMatL']) # 3*3
        self.DistCoeffL = np.array(json_dict['DistCoeffL']) # 1*5
        self.CameraMatR = np.array(json_dict['CameraMatR']) # 3*3
        self.DistCoeffR = np.array(json_dict['DistCoeffR']) # 1*5

        self.R = np.array(json_dict['R'])
        self.T = np.array(json_dict['T'])
        self.E = np.array(json_dict['E'])
        self.F = np.array(json_dict['F'])


    # 旋转向量和平移向量求解
    def calibration_RT(self, points_3D, points_2D, cameraMatrix, distCoeffs):
        points_3D = np.array(points_3D)
        points_2D = np.array(points_2D)
        cameraMatrix = np.array(cameraMatrix).astype(np.float32)
        distCoeffs = np.array(distCoeffs).astype(np.float32)
        _, rvecs, tvecs, inliers = cv2.solvePnPRansac(points_3D.reshape(-1, 1, 3),
                                                    points_2D.reshape(-1, 1, 2),
                                                    cameraMatrix,
                                                    distCoeffs
                                                    )
        R, _ = cv2.Rodrigues(rvecs)
        print('R:\n', R)
        print('rvecs:\n', rvecs)
        print('tvecs:\n', tvecs)

        return R, tvecs
if __name__ == '__main__':
    
    cal_data = Calibration(r"D:\data\xishajietou\weihai\1210")

